# 对角线遍历：https://leetcode-cn.com/problems/diagonal-traverse/

class Solution:
    def findDiagonalOrder(self, mat):
        """ 
            斜上走： row - 1, col + 1. 结束条件(出边界)： row < 0 or col > n - 1
            斜下走： row + 1, col - 1. 结束条件(出边界)： row > m -1 or col < 0  

            结束条件（程序结束）： 遍历完成所有节点 n * m
        """
        # m行 n列
        m, n = len(mat), len(mat[0])
        # 定义动作， 第一个元素是斜向上， 第二个是斜向下
        action = [[-1, 1], [1, -1]]
        rst = []
        row_cur = col_cur = 0
        # 当前动作：斜向上
        act = 0
        i = 0
        while i < n*m:

            # 添加到结果列表
            rst.append(mat[row_cur][col_cur])
            i += 1

            # 计算下次是否改变动作
            # 当前是斜向上
            if act == 0 and (row_cur == 0 or col_cur == n - 1):
                if col_cur == n-1:
                    row_cur += 1
                elif row_cur == 0:
                    col_cur += 1
                act = (act + 1) % 2
                continue
            # 当前时斜向下
            elif act == 1 and (row_cur == m -1 or col_cur == 0):
                if row_cur == m - 1:
                    col_cur += 1
                elif col_cur == 0:
                    row_cur += 1
                act = (act + 1) % 2
                continue

            # 计算元素坐标
            row_cur = row_cur + action[act][0]
            col_cur = col_cur + action[act][1]
        
        return rst

matrix = [
 [ 1, 2, 3 ],
 [ 4, 5, 6 ],
 [ 7, 8, 9 ]
]

s = Solution()
rst = s.findDiagonalOrder(matrix)
print(rst)
